import os
import pandas as pd
import random
from datetime import datetime, timedelta

# 读取股票代码和名称
def read_stock_codes(file_path):
    stock_dict = {}
    with open(file_path, 'r', encoding='utf-8') as file:
        for line in file:
            parts = line.strip().split('\t')
            if len(parts) >= 2:
                stock_dict[parts[0]] = parts[1]
    return stock_dict

# 生成交易日期，避开周末和节假日
def generate_trading_date(start_date, end_date, non_trading_days):
    # 确保开始日期小于结束日期
    if start_date >= end_date:
        start_date = end_date - timedelta(days=30)
    
    max_attempts = 100
    for _ in range(max_attempts):
        random_days = random.randint(0, (end_date - start_date).days)
        date = start_date + timedelta(days=random_days)
        if date.strftime('%Y-%m-%d') not in non_trading_days:
            return date
    
    while start_date <= end_date:
        if start_date.strftime('%Y-%m-%d') not in non_trading_days:
            return start_date
        start_date += timedelta(days=1)
    return end_date

# 生成一只股票的多条交易记录
def generate_stock_trades(stock_code, stock_name, case_id, non_trading_days):
    trades = []
    year = random.choice([2022, 2023, 2024])
    num_buys = random.randint(20, 30)
    
    # 生成第一个交易日期
    start_date = datetime(year, 1, 1)
    end_date = datetime(year, 6, 30)  # 将初始范围限制在上半年
    first_trade_date = generate_trading_date(start_date, end_date, non_trading_days)
    
    # 从第一个交易日期开始，连续生成交易记录
    current_date = first_trade_date
    for i in range(num_buys):
        # 生成买入记录
        quantity = random.randint(500, 2000)
        price = round(random.uniform(5.0, 15.0), 2)
        trades.append({
            "股票代码": stock_code,
            "股票名称": stock_name,
            "成交日期": current_date.strftime('%Y-%m-%d'),
            "买卖方向": "buy",
            "成交数量": quantity,
            "成交价格": price,
            "成交额": round(quantity * price, 2)
        })
        
        # 生成对应的卖出记录
        sell_start = current_date + timedelta(days=1)
        sell_end = min(current_date + timedelta(days=15), datetime(year, 12, 31))
        
        if sell_start < sell_end:
            sell_date = generate_trading_date(sell_start, sell_end, non_trading_days)
            quantity = random.randint(500, 2000)
            price = round(random.uniform(5.0, 15.0), 2)
            trades.append({
                "股票代码": stock_code,
                "股票名称": stock_name,
                "成交日期": sell_date.strftime('%Y-%m-%d'),
                "买卖方向": "sell",
                "成交数量": quantity,
                "成交价格": price,
                "成交额": round(quantity * price, 2)
            })
            
            # 更新下一次买入的日期
            current_date = sell_date + timedelta(days=random.randint(5, 15))
            # 确保不超过年底
            if current_date.year > year:
                break
    
    return trades

def main():
    # 节假日字典
    holidays = {
        2022: ['2022-01-01', '2022-01-31', '2022-04-03', '2022-04-30', 
               '2022-06-03', '2022-09-10', '2022-10-01'],
        2023: ['2023-01-01', '2023-01-21', '2023-04-05', '2023-05-01', 
               '2023-06-22', '2023-09-29', '2023-10-01'],
        2024: ['2024-01-01', '2024-02-10', '2024-04-04', '2024-05-01', 
               '2024-06-10', '2024-09-15', '2024-10-01']
    }

    # 生成周末和节假日日期列表
    non_trading_days = set()
    for year in [2022, 2023, 2024]:
        start_date = datetime(year, 1, 1)
        end_date = datetime(year, 12, 31)
        delta = timedelta(days=1)
        while start_date <= end_date:
            if start_date.weekday() >= 5:
                non_trading_days.add(start_date.strftime('%Y-%m-%d'))
            start_date += delta
        non_trading_days.update(holidays[year])

    # 读取股票代码和名称
    file_path = os.path.abspath('./test/stocks.txt')
    stock_dict = read_stock_codes(file_path)
    
    # 确保有足够的不重复股票可供选择
    if len(stock_dict) < 1000:
        print(f"警告：股票代码数量不足1000个，当前只有{len(stock_dict)}个")
    
    # 随机选择100个不重复的股票
    selected_stocks = random.sample(list(stock_dict.items()), min(1000, len(stock_dict)))
    
    # 生成交易数据
    all_trades = []
    for i, (stock_code, stock_name) in enumerate(selected_stocks, 1):
        trades = generate_stock_trades(stock_code, stock_name, i, non_trading_days)
        all_trades.extend(trades)
    
# 创建DataFrame并按日期排序
    df = pd.DataFrame(all_trades)
    df = df.sort_values(['股票代码', '成交日期'])
    
    # 设置输出文件路径
    output_file = "./test/1000_stocks_test_data.xlsx"
    
    # 如果文件存在则删除
    if os.path.exists(output_file):
        try:
            os.remove(output_file)
            print(f"已删除旧的Excel文件: {output_file}")
        except Exception as e:
            print(f"删除旧文件失败: {str(e)}")
    
    # 保存为新的Excel文件
    df.to_excel(output_file, index=False)
    print(f"生成的Excel文件已保存为: {output_file}")
    
    # 打印统计信息
    unique_stocks = df['股票代码'].nunique()
    print(f"生成的数据包含 {unique_stocks} 个不同的股票")
    print(f"总共生成 {len(df)} 条交易记录")

if __name__ == "__main__":
    main()